{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23cf47df",
   "metadata": {},
   "source": [
    "### Imports and loading "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4980bebf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/salahalzubi/cursor_projects/ROMA-DSPy/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "\u001b[32m2025-11-03 18:30:11.432\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m57\u001b[0m - \u001b[34m\u001b[1mLoading config: path=None, profile=test, overrides=None, env_prefix=ROMA_\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.432\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m66\u001b[0m - \u001b[34m\u001b[1mInitialized empty base config (defaults applied in validation)\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.437\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36m_load_yaml\u001b[0m:\u001b[36m129\u001b[0m - \u001b[34m\u001b[1mLoaded and cached config from config/defaults/config.yaml\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.438\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m81\u001b[0m - \u001b[34m\u001b[1mMerged default config from config/defaults/config.yaml\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.443\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36m_load_yaml\u001b[0m:\u001b[36m129\u001b[0m - \u001b[34m\u001b[1mLoaded and cached config from config/profiles/test.yaml\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.452\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m87\u001b[0m - \u001b[34m\u001b[1mApplied profile: test\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.459\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m103\u001b[0m - \u001b[34m\u001b[1mResolved interpolations\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.460\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: FileToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.460\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: CalculatorToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.460\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: SerperToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.460\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: WebSearchToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.461\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: E2BToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.461\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: BinanceToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.461\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: CoinGeckoToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.461\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: DefiLlamaToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.461\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: ArkhamToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.462\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: CoinglassToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.660\u001b[0m | \u001b[34m\u001b[1mDEBUG   \u001b[0m | \u001b[36mroma_dspy.tools.base.manager\u001b[0m:\u001b[36m_register_toolkit_class\u001b[0m:\u001b[36m127\u001b[0m - \u001b[34m\u001b[1mRegistered toolkit: MCPToolkit\u001b[0m\n",
      "\u001b[32m2025-11-03 18:30:11.660\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mroma_dspy.config.manager\u001b[0m:\u001b[36mload_config\u001b[0m:\u001b[36m110\u001b[0m - \u001b[1mConfiguration loaded and validated successfully\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "# Add project root to Python path\n",
    "project_root = Path.cwd().parent\n",
    "if str(project_root) not in sys.path:\n",
    "    sys.path.insert(0, str(project_root))\n",
    "\n",
    "import asyncio\n",
    "import dspy\n",
    "from dspy import GEPA\n",
    "\n",
    "from prompt_optimization.config import (\n",
    "    LMConfig,\n",
    ")\n",
    "from prompt_optimization.dataset_loaders import (\n",
    "    load_aimo_datasets,\n",
    "    load_frames_dataset,\n",
    "    load_seal0_dataset,\n",
    "    load_simpleqa_verified_dataset,\n",
    ")\n",
    "from prompt_optimization.judge import ComponentJudge\n",
    "from prompt_optimization.metrics import MetricWithFeedback\n",
    "from prompt_optimization.optimizer import create_optimizer\n",
    "from prompt_optimization.prompts import (\n",
    "    GRADER_PROMPT,\n",
    ")\n",
    "from roma_dspy.config import load_config\n",
    "from roma_dspy.core.engine.solve import RecursiveSolver\n",
    "from roma_dspy.core.modules.recursive_solver import RecursiveSolverModule\n",
    "\n",
    "dspy.settings.provide_traceback = True  # optional but mirrors the old notebook\n",
    "import os\n",
    "os.chdir(project_root)  # Change to project root for config loading\n",
    "\n",
    "opt_cfg = load_config(profile=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa039eb3",
   "metadata": {},
   "source": [
    "### Config LLMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7f020b88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch the knobs you used to tweak in the notebook.\n",
    "opt_cfg.train_size = 32\n",
    "opt_cfg.val_size = 12\n",
    "opt_cfg.test_size = 12\n",
    "opt_cfg.dataset_seed = 42\n",
    "opt_cfg.max_metric_calls = 225\n",
    "opt_cfg.max_depth = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6af57c",
   "metadata": {},
   "source": [
    "### Init solvers and what not"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "700e9d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = RecursiveSolver(\n",
    "    config=opt_cfg,\n",
    "    max_depth=opt_cfg.max_depth,\n",
    "    enable_logging=opt_cfg.enable_logging,\n",
    "    enable_checkpoints=False,\n",
    ")\n",
    "solver_module = RecursiveSolverModule(solver=solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d6081f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from prompt_optimization.metrics.search_metric import SearchMetric\n",
    "from prompt_optimization.prompts.grader_prompts.search_grader_prompt import SEARCH_GRADER_PROMPT\n",
    "\n",
    "\n",
    "scorer_lm = LMConfig(\"openrouter/google/gemini-2.5-flash-lite-preview-09-2025\", temperature=0.3, max_tokens=128000, cache=True)\n",
    "component_lm = LMConfig(\"openrouter/anthropic/claude-sonnet-4.5\", temperature=0.45, max_tokens=128000, cache=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "492816a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_metric_scorer = SearchMetric(prompt=SEARCH_GRADER_PROMPT, lm_config=scorer_lm)\n",
    "component_metric_scorer = ComponentJudge(prompt=GRADER_PROMPT, lm_config=component_lm)\n",
    "metric = MetricWithFeedback(judge = component_metric_scorer, scoring_metric = search_metric_scorer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83b667a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_set, val_set, test_set = load_frames_dataset(\n",
    "    train_size=opt_cfg.train_size,\n",
    "    val_size=opt_cfg.val_size,\n",
    "    test_size=opt_cfg.test_size,\n",
    "    seed=opt_cfg.dataset_seed,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d25b0d72",
   "metadata": {},
   "source": [
    "### Perform an eval on the val set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcacaefc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# executor = AsyncParallelExecutor(max_concurrency=4)\n",
    "\n",
    "# results = await executor.execute_batch(solver_module, test_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d486a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc320b30",
   "metadata": {},
   "source": [
    "### Prompt tuning stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e34080e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = GEPA(\n",
    "    metric=metric,\n",
    "    # auto=\"light\",\n",
    "    component_selector=\"round_robin\",\n",
    "    max_metric_calls=48,\n",
    "    add_format_failure_as_feedback=True,\n",
    "    num_threads=8,\n",
    "    track_stats=True,\n",
    "    log_dir=\"logs/frames_test\",\n",
    "    # use_wandb=True,\n",
    "    # wandb_init_kwargs={\"project\": \"aime_test\"},\n",
    "    reflection_minibatch_size=8,\n",
    "    reflection_lm=dspy.LM(model=\"openrouter/anthropic/claude-sonnet-4.5\", temperature=.75, max_tokens=128000)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85bcee5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_program = optimizer.compile(\n",
    "    solver_module,\n",
    "    trainset=train_set[:8],\n",
    "    valset=val_set,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f25543c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dir(optimized_program.named_predictors()[1][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43948e36",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
